{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ur8xi4C7S06n"
      },
      "outputs": [],
      "source": [
        "# Copyright 2020 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1c955a52adac"
      },
      "source": [
        "# Feedback or issues?\n",
        "For any feedback or questions, please open an [issue](https://github.com/googleapis/python-aiplatform/issues)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eHLV0D7Y5jtU"
      },
      "source": [
        "# Vertex SDK for Python: Custom Training Example with Unmanaged Image Dataset\n",
        "\n",
        "\n",
        "To use this Colaboratory notebook, you copy the notebook to your own Google Drive and open it with Colaboratory (or Colab). You can run each step, or cell, and see its results. To run a cell, use Shift+Enter. Colab automatically displays the return value of the last line in each cell. For more information about running notebooks in Colab, see the [Colab welcome page](https://colab.research.google.com/notebooks/welcome.ipynb).\n",
        "\n",
        "This notebook demonstrate how to create a custom model based on an image dataset. It will require you provide a bucket where the dataset will be stored.\n",
        "\n",
        "Note: you may incur charges for training, prediction,  storage or usage of other GCP products in connection with testing this SDK."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lld3eeJUs5yM"
      },
      "source": [
        "# Install Vertex SDK for Python, Authenticate, and upload of a Dataset to your GCS bucket\n",
        "\n",
        "\n",
        "After the SDK installation the kernel will be automatically restarted. You may see this error message `Your session crashed for an unknown reason` which is normal."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sBfZtR4X1Dr_"
      },
      "outputs": [],
      "source": [
        "!pip3 uninstall -y google-cloud-aiplatform\n",
        "!pip3 install google-cloud-aiplatform\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c0SNmTBeD2nV"
      },
      "source": [
        "### Enter your project and GCS bucket\n",
        "\n",
        "Enter your Project Id in the cell below. Then run the cell to make sure the Cloud SDK uses the right project for all the commands in this notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f5be5dce7259"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kng9iKBwqcS5"
      },
      "outputs": [],
      "source": [
        "MY_PROJECT = \"YOUR PROJECT\"\n",
        "MY_STAGING_BUCKET = \"gs://YOUR BUCKET\"  # bucket should be in same region as ucaip"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rk43VP_IqcTE"
      },
      "source": [
        "# Initialize Vertex SDK for Python\n",
        "\n",
        "Initialize the *client* for Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VCiC9gBWqcTF"
      },
      "outputs": [],
      "source": [
        "from google.cloud import aiplatform\n",
        "\n",
        "aiplatform.init(project=MY_PROJECT, staging_bucket=MY_STAGING_BUCKET)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VcEOYYolqcTN"
      },
      "source": [
        "# Write your Training Script\n",
        "- Write this cell as a file which will be used for custom training.\n",
        "- Instead of using a managed dataset, a Tensorflow Dataset URI is passed in through the 'args' parameter of the 'run' function. The script will download the data from the URI at training time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78SHpd0tt8UQ"
      },
      "outputs": [],
      "source": [
        "%%writefile training_script.py\n",
        "\n",
        "# Source: https://cloud.google.com/vertex-ai/docs/tutorials/image-recognition-custom\n",
        "\n",
        "import argparse\n",
        "import logging\n",
        "import os\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_datasets as tfds\n",
        "\n",
        "IMG_WIDTH = 128\n",
        "\n",
        "def normalize_img(image):\n",
        "    \"\"\"Normalizes image.\n",
        "\n",
        "    * Resizes image to IMG_WIDTH x IMG_WIDTH pixels\n",
        "    * Casts values from `uint8` to `float32`\n",
        "    * Scales values from [0, 255] to [0, 1]\n",
        "\n",
        "    Returns:\n",
        "      A tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color channels)\n",
        "    \"\"\"\n",
        "    image = tf.image.resize_with_pad(image, IMG_WIDTH, IMG_WIDTH)\n",
        "    return image / 255.\n",
        "\n",
        "\n",
        "def normalize_img_and_label(image, label):\n",
        "    \"\"\"Normalizes image and label.\n",
        "\n",
        "    * Performs normalize_img on image\n",
        "    * Passes through label unchanged\n",
        "\n",
        "    Returns:\n",
        "      Tuple (image, label) where\n",
        "      * image is a tensor with shape (IMG_WIDTH, IMG_WIDTH, 3). (3 color\n",
        "        channels)\n",
        "      * label is an unchanged integer [0, 4] representing flower type\n",
        "    \"\"\"\n",
        "    return normalize_img(image), label\n",
        "\n",
        "def get_args():\n",
        "  \"\"\"Argument parser.\n",
        "  Returns:\n",
        "    Dictionary of arguments.\n",
        "  \"\"\"\n",
        "  parser = argparse.ArgumentParser(description='Flower classification sample')\n",
        "  parser.add_argument(\n",
        "      '--tfds',\n",
        "      default=None,\n",
        "      help='The tfds URI from https://www.tensorflow.org/datasets/ to load the data from')\n",
        "\n",
        "  args = parser.parse_args()\n",
        "  return args\n",
        "\n",
        "# Training settings\n",
        "args = get_args()\n",
        "\n",
        "if 'AIP_MODEL_DIR' not in os.environ:\n",
        "    raise KeyError(\n",
        "        'The `AIP_MODEL_DIR` environment variable has not been' +\n",
        "        'set. See https://cloud.google.com/vertex-ai/docs/tutorials/image-recognition-custom/training'\n",
        "    )\n",
        "output_directory = os.environ['AIP_MODEL_DIR']\n",
        "\n",
        "logging.info('Loading and preprocessing data ...')\n",
        "dataset = tfds.load(args.tfds,\n",
        "                    split='train',\n",
        "                    try_gcs=True,\n",
        "                    shuffle_files=True,\n",
        "                    as_supervised=True)\n",
        "dataset = dataset.map(normalize_img_and_label,\n",
        "                      num_parallel_calls=tf.data.experimental.AUTOTUNE)\n",
        "dataset = dataset.cache()\n",
        "dataset = dataset.shuffle(1000)\n",
        "dataset = dataset.batch(128)\n",
        "dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)\n",
        "\n",
        "logging.info('Creating and training model ...')\n",
        "model = tf.keras.Sequential([\n",
        "    tf.keras.layers.Conv2D(16,\n",
        "                           3,\n",
        "                           padding='same',\n",
        "                           activation='relu',\n",
        "                           input_shape=(IMG_WIDTH, IMG_WIDTH, 3)),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Conv2D(32, 3, padding='same', activation='relu'),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu'),\n",
        "    tf.keras.layers.MaxPooling2D(),\n",
        "    tf.keras.layers.Flatten(),\n",
        "    tf.keras.layers.Dense(512, activation=\"relu\"),\n",
        "    tf.keras.layers.Dense(5)  # 5 classes\n",
        "])\n",
        "model.compile(\n",
        "    optimizer='adam',\n",
        "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
        "    metrics=['accuracy'])\n",
        "model.fit(dataset, epochs=10)\n",
        "\n",
        "logging.info(f'Exporting SavedModel to: {output_directory}')\n",
        "# Add softmax layer for intepretability\n",
        "probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])\n",
        "probability_model.save(output_directory)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6-bBqipfqcTS"
      },
      "source": [
        "# Launch a Training Job to Create a Model\n",
        "\n",
        "Once we have defined your training script, we will create a model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "btb6d48lqcTT"
      },
      "outputs": [],
      "source": [
        "job = aiplatform.CustomTrainingJob(\n",
        "    display_name=\"train-flowers-dist-1-replica\",\n",
        "    script_path=\"training_script.py\",\n",
        "    container_uri=\"gcr.io/cloud-aiplatform/training/tf-cpu.2-2:latest\",\n",
        "    requirements=[\"gcsfs==0.7.1\"],\n",
        "    model_serving_container_image_uri=\"gcr.io/cloud-aiplatform/prediction/tf2-cpu.2-2:latest\",\n",
        ")\n",
        "model = job.run(\n",
        "    args=[\"--tfds\", \"tf_flowers:3.*.*\"],\n",
        "    replica_count=1,\n",
        "    model_display_name=\"flowers-model\",\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5vhDsMJNqcTW"
      },
      "source": [
        "# Deploy Your Model\n",
        "\n",
        "Deploy your model, then wait until the model FINISHES deployment before proceeding to prediction.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Y9GH72wWqcTX"
      },
      "outputs": [],
      "source": [
        "endpoint = model.deploy(machine_type=\"n1-standard-4\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIw1ifPuqcTb"
      },
      "source": [
        "# Predict on the Endpoint\n",
        "To do a prediction you will need some flowers images. You can download some photos of flowers or use the ones provided below."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wn8fUtq389r4"
      },
      "outputs": [],
      "source": [
        "!gsutil -m cp -R gs://cloud-ml-data/img/flower_photos/daisy/14221848160_7f0a37c395.jpg .\n",
        "!gsutil -m cp -R gs://cloud-ml-data/img/flower_photos/tulips/13289268363_b9337d751e.jpg .\n",
        "!gsutil -m cp -R gs://cloud-ml-data/img/flower_photos/sunflowers/14623719696_1bb7970208_n.jpg ."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ktakP9r7mCt"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from PIL import Image\n",
        "\n",
        "daisy_floats = np.array(Image.open(\"14221848160_7f0a37c395.jpg\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xzGidweY8Dgy"
      },
      "outputs": [],
      "source": [
        "small_image = np.array(Image.fromarray(np.uint8(daisy_floats)).resize((128, 128)))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cL6VdVufgH8a"
      },
      "outputs": [],
      "source": [
        "endpoint.predict(instances=[small_image.tolist()])"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "AI_Platform_(Unified)_SDK_Custom_Training_with_Unmanaged_Image_Dataset.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
